"""
多模态数据加载器模块

该模块提供了加载和处理多模态数据集的功能，支持图像和文本数据的联合处理。
主要功能包括：
- 加载JSONL格式的多模态数据集
- 自动处理图像路径和文本内容
- 支持数据预处理和验证
- 提供灵活的数据过滤和采样功能

使用方法：
    from common.multimodal_data_loader import MultimodalDataLoader
    loader = MultimodalDataLoader(config)
    dataset = loader.load_dataset()
"""

import json
import os
from typing import List, Dict, Any, Optional, Union
import logging
from pathlib import Path

# 设置日志
logger = logging.getLogger(__name__)


class MultimodalDataLoader:
    """
    多模态数据加载器
    
    该类负责加载和处理多模态数据集，支持图像和文本数据的联合处理。
    能够自动识别数据格式，处理图像路径，并提供数据验证功能。
    
    Attributes:
        config: 配置对象，包含数据相关的配置参数
        image_root (str): 图像文件的根目录路径
        dataset_path (str): 数据集文件路径
        max_samples (int): 最大样本数量，-1表示加载所有样本
    """
    
    def __init__(self, config):
        """
        初始化多模态数据加载器
        
        Args:
            config: 配置对象，包含数据相关的配置参数
        """
        self.config = config
        self.image_root = config.data.image_root
        self.dataset_path = config.data.dataset_path
        self.max_samples = config.data.max_samples
        
        # 获取字段映射配置，默认为标准字段
        self.field_mapping = config.data.field_mapping or {
            'question': 'question',
            'answer': 'answer', 
            'image': 'image'
        }
        
        logger.info(f"初始化多模态数据加载器")
        logger.info(f"数据集路径: {self.dataset_path}")
        logger.info(f"图像根目录: {self.image_root}")
        logger.info(f"最大样本数: {self.max_samples}")
        logger.info(f"字段映射: {self.field_mapping}")
    
    def load_dataset(self) -> List[Dict[str, Any]]:
        """
        加载数据集的主方法
        
        从指定路径加载数据集文件，支持JSON和JSONL格式。
        自动处理图像路径，验证数据完整性，并应用采样限制。
        
        Returns:
            List[Dict[str, Any]]: 处理后的数据集列表
            
        Raises:
            FileNotFoundError: 数据集文件不存在
            ValueError: 数据格式错误
            
        Example:
            >>> loader = MultimodalDataLoader(config)
            >>> dataset = loader.load_dataset()
            >>> print(f"加载了 {len(dataset)} 个样本")
        """
        # 检查数据集文件是否存在
        if not os.path.exists(self.dataset_path):
            raise FileNotFoundError(f"数据集文件不存在: {self.dataset_path}")
        
        logger.info(f"开始加载数据集: {self.dataset_path}")
        
        # 根据文件扩展名选择加载方式
        if self.dataset_path.endswith('.jsonl'):
            raw_data = self._load_jsonl(self.dataset_path)
        elif self.dataset_path.endswith('.json'):
            raw_data = self._load_json(self.dataset_path)
        else:
            raise ValueError(f"不支持的数据格式: {self.dataset_path}")
        
        logger.info(f"原始数据加载完成，样本数量: {len(raw_data)}")
        
        # 处理数据
        processed_data = []
        for i, item in enumerate(raw_data):
            try:
                processed_item = self._process_item(item, i)
                processed_data.append(processed_item)
                
                # 注释掉max_samples限制，让数据加载器加载所有数据
                # 数据范围选择将在主程序中处理
                # if self.max_samples > 0 and len(processed_data) >= self.max_samples:
                #     logger.info(f"达到最大样本数限制: {self.max_samples}")
                #     break
                    
            except Exception as e:
                logger.warning(f"处理样本 {i} 时出错: {e}")
                continue
        
        logger.info(f"数据处理完成，有效样本数量: {len(processed_data)}")
        
        # 验证数据集
        self._validate_dataset(processed_data)
        
        return processed_data
    
    def _load_jsonl(self, file_path: str) -> List[Dict[str, Any]]:
        """
        加载JSONL格式文件
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            List[Dict[str, Any]]: 数据列表
        """
        data = []
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if line:  # 跳过空行
                        try:
                            item = json.loads(line)
                            data.append(item)
                        except json.JSONDecodeError as e:
                            logger.warning(f"JSONL文件第 {line_num} 行格式错误: {e}")
                            continue
        except Exception as e:
            raise ValueError(f"读取JSONL文件失败: {e}")
        
        return data
    
    def _load_json(self, file_path: str) -> List[Dict[str, Any]]:
        """
        加载JSON格式文件
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            List[Dict[str, Any]]: 数据列表
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                
            # 确保数据是列表格式
            if not isinstance(data, list):
                raise ValueError("JSON文件应包含数据列表")
                
            return data
            
        except Exception as e:
            raise ValueError(f"读取JSON文件失败: {e}")
    
    def _process_item(self, item: Dict[str, Any], index: int) -> Dict[str, Any]:
        """
        处理单个数据项
        
        从原始数据项中提取问题、答案和图像信息，并进行标准化处理。
        
        Args:
            item (Dict[str, Any]): 原始数据项
            index (int): 数据项索引
            
        Returns:
            Dict[str, Any]: 处理后的数据项
            
        Raises:
            ValueError: 数据项格式错误
        """
        # 提取问题文本
        question = self._extract_question(item)
        if not question:
            raise ValueError(f"样本 {index} 缺少问题文本")
        
        # 提取答案
        answer = self._extract_answer(item)
        
        # 提取图像路径
        image_path = self._extract_image_path(item)
        
        # 处理图像路径
        processed_image_path = None
        has_image = False
        
        if image_path:
            processed_image_path = self._process_image_path(image_path)
            has_image = processed_image_path is not None
        
        # 提取原始ID（如果存在）
        original_id = item.get('id', str(index))
        
        # 构建处理后的数据项
        processed_item = {
            'id': original_id,  # 保留原始ID
            'sample_id': index,  # 保留索引用于兼容性
            'question': question,
            'answer': answer,
            'image_path': processed_image_path,
            'has_image': has_image,
            'original_data': item
        }
        
        return processed_item
    
    def _extract_question(self, item: Dict[str, Any]) -> str:
        """
        从数据项中提取问题文本
        
        根据配置的字段映射提取问题，支持回退到常见字段名称。
        
        Args:
            item (Dict[str, Any]): 数据项
            
        Returns:
            str: 问题文本
        """
        # 首先尝试使用配置的字段映射
        question_field = self.field_mapping.get('question', 'question')
        if question_field in item and item[question_field]:
            return str(item[question_field]).strip()
        
        # 回退到常见的问题字段名称
        fallback_fields = ['question', 'instruction', 'prompt', 'input', 'query']
        for field in fallback_fields:
            if field in item and item[field]:
                return str(item[field]).strip()
        
        return ""
    
    def _extract_answer(self, item: Dict[str, Any]) -> str:
        """
        从数据项中提取答案文本
        
        根据配置的字段映射提取答案，支持回退到常见字段名称。
        
        Args:
            item (Dict[str, Any]): 数据项
            
        Returns:
            str: 答案文本
        """
        # 首先尝试使用配置的字段映射
        answer_field = self.field_mapping.get('answer', 'answer')
        if answer_field in item and item[answer_field] is not None:
            return str(item[answer_field]).strip()
        
        # 回退到常见的答案字段名称
        fallback_fields = ['answer', 'output', 'response', 'target', 'label']
        for field in fallback_fields:
            if field in item and item[field] is not None:
                return str(item[field]).strip()
        
        return ""
    
    def _extract_image_path(self, item: Dict[str, Any]) -> Optional[str]:
        """
        从数据项中提取图像路径
        
        根据配置的字段映射提取图像路径，支持回退到常见字段名称。
        
        Args:
            item (Dict[str, Any]): 数据项
            
        Returns:
            Optional[str]: 图像路径
        """
        # 首先尝试使用配置的字段映射
        image_field = self.field_mapping.get('image', 'image')
        if image_field in item and item[image_field]:
            return str(item[image_field]).strip()
        
        # 回退到常见的图像路径字段名称
        fallback_fields = ['image_path', 'image', 'img_path', 'img', 'image_file']
        for field in fallback_fields:
            if field in item and item[field]:
                return str(item[field]).strip()
        
        return None
    
    def _process_image_path(self, image_path: str) -> Optional[str]:
        """
        处理图像路径
        
        将相对路径转换为绝对路径，并验证文件是否存在。
        
        Args:
            image_path (str): 原始图像路径
            
        Returns:
            Optional[str]: 处理后的图像路径，如果文件不存在则返回None
        """
        # 如果已经是绝对路径，直接使用
        if os.path.isabs(image_path):
            full_path = image_path
        else:
            # 相对路径，与image_root拼接
            full_path = os.path.join(self.image_root, image_path)
        
        # 标准化路径
        full_path = os.path.normpath(full_path)
        
        # 检查文件是否存在
        if os.path.exists(full_path):
            return full_path
        else:
            logger.warning(f"图像文件不存在: {full_path}")
            return None
    
    def _validate_dataset(self, dataset: List[Dict[str, Any]]) -> None:
        """
        验证数据集的完整性
        
        检查数据集的基本统计信息和数据质量。
        
        Args:
            dataset (List[Dict[str, Any]]): 数据集
        """
        if not dataset:
            raise ValueError("数据集为空")
        
        # 统计信息
        total_samples = len(dataset)
        image_samples = sum(1 for item in dataset if item['has_image'])
        text_only_samples = total_samples - image_samples
        
        # 检查问题和答案的完整性
        missing_questions = sum(1 for item in dataset if not item['question'])
        missing_answers = sum(1 for item in dataset if not item['answer'])
        
        logger.info("=== 数据集验证结果 ===")
        logger.info(f"总样本数: {total_samples}")
        logger.info(f"图像样本数: {image_samples}")
        logger.info(f"纯文本样本数: {text_only_samples}")
        logger.info(f"缺少问题的样本数: {missing_questions}")
        logger.info(f"缺少答案的样本数: {missing_answers}")
        
        # 警告信息
        if missing_questions > 0:
            logger.warning(f"有 {missing_questions} 个样本缺少问题文本")
        
        if missing_answers > 0:
            logger.warning(f"有 {missing_answers} 个样本缺少答案文本")
        
        if image_samples == 0:
            logger.warning("数据集中没有图像样本")
    
    def get_sample_by_id(self, dataset: List[Dict[str, Any]], sample_id: int) -> Optional[Dict[str, Any]]:
        """
        根据ID获取样本
        
        Args:
            dataset (List[Dict[str, Any]]): 数据集
            sample_id (int): 样本ID
            
        Returns:
            Optional[Dict[str, Any]]: 样本数据，如果不存在则返回None
        """
        for item in dataset:
            if item['sample_id'] == sample_id:
                return item
        return None
    
    def filter_by_type(self, dataset: List[Dict[str, Any]], has_image: bool) -> List[Dict[str, Any]]:
        """
        根据类型过滤数据集
        
        Args:
            dataset (List[Dict[str, Any]]): 数据集
            has_image (bool): 是否包含图像
            
        Returns:
            List[Dict[str, Any]]: 过滤后的数据集
        """
        return [item for item in dataset if item['has_image'] == has_image]
    
    def get_statistics(self, dataset: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        获取数据集统计信息
        
        Args:
            dataset (List[Dict[str, Any]]): 数据集
            
        Returns:
            Dict[str, Any]: 统计信息
        """
        total_samples = len(dataset)
        image_samples = sum(1 for item in dataset if item['has_image'])
        text_only_samples = total_samples - image_samples
        
        # 计算问题和答案的平均长度
        question_lengths = [len(item['question']) for item in dataset if item['question']]
        answer_lengths = [len(item['answer']) for item in dataset if item['answer']]
        
        avg_question_length = sum(question_lengths) / len(question_lengths) if question_lengths else 0
        avg_answer_length = sum(answer_lengths) / len(answer_lengths) if answer_lengths else 0
        
        return {
            'total_samples': total_samples,
            'image_samples': image_samples,
            'text_only_samples': text_only_samples,
            'image_ratio': image_samples / total_samples if total_samples > 0 else 0,
            'avg_question_length': avg_question_length,
            'avg_answer_length': avg_answer_length,
            'missing_questions': sum(1 for item in dataset if not item['question']),
            'missing_answers': sum(1 for item in dataset if not item['answer'])
        }